from utils.vit_utils import get_classes
from vit_pytorch.vit import ViTBody

import torch
import torch.nn as nn


class ViT(object):
    _defaults = {
        # --------------------------------------------------------------------------#
        #   使用自己训练好的模型进行预测一定要修改model_path和classes_path！
        #   model_path指向logs文件夹下的权值文件，classes_path指向model_data下的txt
        #
        #   训练好后logs文件夹下存在多个权值文件，选择验证集损失较低的即可。
        #   验证集损失较低不代表mAP较高，仅代表该权值在验证集上泛化性能较好。
        #   如果出现shape不匹配，同时要注意训练时的model_path和classes_path参数的修改
        # --------------------------------------------------------------------------#
        "model_path": "pretrain_weights/imagenet21k+imagenet2012_ViT-L_32.npz",
        "classes_path": 'model_data/classes.txt',
        # ---------------------------------------------------------------------#
        #   输入图片的大小，必须为32的倍数。
        # ---------------------------------------------------------------------#
        "image_size": 256,
        # ---------------------------------------------------------------------#
        #   只有得分大于置信度的预测框会被保留下来
        # ---------------------------------------------------------------------#
        "patch_size": 32,
        # -------------------------------#
        #   是否使用Cuda
        #   没有GPU可以设置成False
        # -------------------------------#
        "cuda": True,
        # -------------------------------#
        #   设置dropout
        # -------------------------------#
        "dropout": 0.1,
        # -------------------------------#
        #   设置emb_dropout
        # -------------------------------#
        "emb_dropout": 0.1,
    }

    @classmethod
    def get_defaults(cls, n):
        if n in cls._defaults:
            return cls._defaults[n]
        else:
            return "Unrecognized attribute name '" + n + "'"

    # ---------------------------------------------------#
    #   初始化VIT
    # ---------------------------------------------------#
    def __init__(self, **kwargs):
        self.__dict__.update(self._defaults)

        # 将配置项：从对象dict中获取  并set为对象属性
        for name, value in kwargs.items():
            setattr(self, name, value)

        # ---------------------------------------------------#
        #   获取分类总数、以及类别
        # ---------------------------------------------------#
        self.class_names, self.num_classes = get_classes(self.classes_path)

        # ---------------------------------------------------#
        #   画框设置不同的颜色
        # ---------------------------------------------------#
        self.generate()

        # show_config(**self._defaults)
        # ---------------------------------------------------#
        #   生成模型
        # ---------------------------------------------------#

    def generate(self, onnx=False):
        # ---------------------------------------------------#
        #   建立ViT模型，载入vit模型的权重
        # ---------------------------------------------------#
        self.net = ViTBody(image_size=self.image_size, patch_size=self.patch_size, num_classes=self.num_classes, dropout=self.dropout, emb_dropout=self.emb_dropout)
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.net.load_state_dict(torch.load(self.model_path, map_location=device))
        self.net = self.net.eval()
        print('{} model, anchors, and classes loaded.'.format(self.model_path))
        if not onnx:
            if self.cuda:
                self.net = nn.DataParallel(self.net).cuda()

    # ---------------------------------------------------#
    #   检测图片
    # ---------------------------------------------------#
    def detect_image(self, image, transform=None):
        # ---------------------------------------------------------#
        #   使用transform处理图片,并添加batch_size 维度
        # ---------------------------------------------------------#
        image_data = transform(image).unsqueeze(0)
        # ---------------------------------------------------------#
        #   转换成tensor
        # ---------------------------------------------------------#
        with torch.no_grad():
            images = torch.from_numpy(image_data)
            if self.cuda:
                images = images.cuda()
            # ---------------------------------------------------------#
            #   将图像输入网络当中进行预测！
            # ---------------------------------------------------------#
            outputs = self.net(images)
            print("predict result is {}", outputs)
            # 将图片通过model正向传播，得到输出，将输入进行压缩，
            # 将batch维度压缩掉，得到最终输出（out）
            # 经过softmax处理后，就变成概率分布的形式了
            # 通过argmax方法，得到概率最大的处所对应的索引
            outputs = torch.squeeze(outputs)
            predict = torch.softmax(outputs, dim=0)
            predict_cla = torch.argmax(predict).numpy()
            cla_ = self.class_names[predict_cla]
        return cla_


